"""GIF 生成器对外 API：单帧渲染与 GIF 合成工具。

此模块把单帧渲染封装为返回 PNG bytes 的函数，便于在子进程中调用或在父进程中复用。
支持两种模式：完整渲染（包含 colorbar 与帧标签）和仅渲染静态背景（用于缓存）。
当传入 background_png 时，会把动态层渲染为透明 PNG 并与 background 合成，返回合成后的 PNG bytes。
"""
from typing import List, Dict, Any
import io
import json
from pathlib import Path

import matplotlib as mpl

# 延迟导入 matplotlib.pyplot 以便在子进程设置 Agg 后调用

def _ensure_agg():
    try:
        mpl.use('Agg')
    except Exception:
        pass


def render_frame_bytes(features: List[Dict[str, Any]], params: Dict[str, Any]) -> bytes:
    """Render a single frame (from GeoJSON-like features) and return PNG bytes.

    支持 params 字段:
    - figsize, dpi, column, cmap, vmin, vmax, colorbar_bbox, font_family, ...
    - render_background_only (bool): 仅渲染静态背景（无 colorbar / 无帧标签）
    - background_png (bytes): 如果提供，则在透明画布上渲染动态层并与背景合成
    """
    _ensure_agg()
    import matplotlib.pyplot as plt
    from shapely.geometry import shape
    from matplotlib.ticker import FuncFormatter
    import pandas as pd
    import geopandas as gpd
    from PIL import Image as PILImage
    # 尝试设置中文字体（若提供）以减少子进程中文缺字警告
    try:
        from utils.plot_style import set_chinese_font, set_plot_style
        set_plot_style(params.get('plot_cfg', {}))
        if params.get('font_family'):
            set_chinese_font(params.get('font_family'), quiet=True)
    except Exception:
        pass

    # 构建 GeoDataFrame
    props = [feat.get('properties', {}) for feat in features]
    geoms = [shape(feat.get('geometry')) for feat in features]
    df = pd.DataFrame(props)
    gdf = gpd.GeoDataFrame(df, geometry=geoms, crs=params.get('crs'))

    # 基本样式与参数
    figsize = params.get('figsize', (10, 12))
    dpi = params.get('dpi', 150)
    column = params.get('column')
    cmap = params.get('cmap', 'Reds')
    vmin = params.get('vmin')
    vmax = params.get('vmax')
    render_bg_only = params.get('render_background_only', False)
    background_png = params.get('background_png')

    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
    ax.set_axis_off()

    # 渲染逻辑分支：背景渲染 vs 动态渲染
    if render_bg_only:
        # 仅绘制地理外壳（不填充动态值）
        gdf.plot(color=params.get('background_fill', '#f7f7f7'), edgecolor=params.get('edgecolor', '#444444'), linewidth=params.get('linewidth', 0.2), ax=ax)
        # 不绘制 colorbar / 标签
        buf = io.BytesIO()
        fig.savefig(buf, format='png', dpi=dpi, bbox_inches='tight')
        plt.close(fig)
        buf.seek(0)
        data = buf.getvalue()
        buf.close()
        return data

    if background_png:
        # 渲染动态层到透明画布，然后与背景合成
        # 在透明画布上绘制 choropleth（不画 colorbar 与标签）
        # 确保 color mapping 使用相同的 vmin/vmax
        norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
        # 使用 alpha=1 在透明背景上绘制
        gdf.plot(column=column, cmap=cmap, linewidth=params.get('linewidth', 0.2), edgecolor='none', ax=ax, norm=norm)
        ax.set_axis_off()
        buf = io.BytesIO()
        # 保存为包含 alpha 的 PNG
        fig.savefig(buf, format='png', dpi=dpi, bbox_inches='tight', transparent=True)
        plt.close(fig)
        buf.seek(0)
        overlay_bytes = buf.getvalue()
        buf.close()

        # 使用 PIL 合成 background + overlay
        bg = PILImage.open(io.BytesIO(background_png)).convert('RGBA')
        overlay = PILImage.open(io.BytesIO(overlay_bytes)).convert('RGBA')
        # 如果尺寸不同，按背景尺寸缩放 overlay
        if overlay.size != bg.size:
            overlay = overlay.resize(bg.size, resample=PILImage.BILINEAR)
        bg.paste(overlay, (0, 0), overlay)
        out_buf = io.BytesIO()
        bg.save(out_buf, format='PNG')
        out_buf.seek(0)
        data = out_buf.getvalue()
        out_buf.close()
        return data

    # 默认：完整绘制（包含 colorbar 与帧标签）
    norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
    gdf.plot(column=column, cmap=cmap, linewidth=params.get('linewidth', 0.2), edgecolor=params.get('edgecolor', '#444444'), ax=ax, norm=norm)
    mappable = mpl.cm.ScalarMappable(norm=norm, cmap=cmap)
    mappable._A = []
    cax = fig.add_axes(params.get('colorbar_bbox', (0.15, 0.16, 0.7, 0.03)))
    cb = fig.colorbar(mappable, cax=cax, orientation='horizontal')

    # colorbar formatter 支持（如果有）
    if params.get('colorbar_tick_percent'):
        from src.gif_generator import _format_cb_ticks_percent, _format_cb_ticks_percent_identity
        if params.get('colorbar_tick_percent_multiply', True):
            cb.ax.xaxis.set_major_formatter(FuncFormatter(_format_cb_ticks_percent))
        else:
            cb.ax.xaxis.set_major_formatter(FuncFormatter(_format_cb_ticks_percent_identity))
    else:
        from src.gif_generator import _format_cb_ticks_identity
        cb.ax.xaxis.set_major_formatter(FuncFormatter(_format_cb_ticks_identity))
    cb.ax.tick_params(labelsize=params.get('colorbar_label_size', 9))
    cb.set_label(params.get('colorbar_label', ''), fontsize=params.get('colorbar_label_size', 11))

    # 帧标签（若提供 frame_label_parts）
    flp = params.get('frame_label_parts') or {}
    x_base = 0.97
    y_base = 0.05
    bbox_kwargs = dict(facecolor='black', alpha=0.35, pad=4)
    year_label_kwargs = dict(color=params.get('frame_label_color', 'white'), fontsize=params.get('frame_label_size', 18), weight='bold', va='bottom')
    if flp.get('month_with_year'):
        year_text = flp.get('year', '')
        month_text = flp.get('month', '')
        y_kwargs = year_label_kwargs.copy()
        m_kwargs = year_label_kwargs.copy()
        y_kwargs['color'] = params.get('frame_label_year_color', y_kwargs.get('color', 'white'))
        m_kwargs['color'] = params.get('frame_label_month_color', '#ffd700')
        ax.text(x_base, y_base, month_text, transform=ax.transAxes, bbox=bbox_kwargs, ha='right', **m_kwargs)
        ax.text(x_base - 0.08, y_base, f"{year_text}年", transform=ax.transAxes, bbox=bbox_kwargs, ha='right', **y_kwargs)
    else:
        single_kwargs = year_label_kwargs.copy()
        single_kwargs['color'] = params.get('frame_label_color', single_kwargs.get('color', 'white'))
        ax.text(x_base, y_base, str(params.get('frame_label_str', '')), transform=ax.transAxes, bbox=bbox_kwargs, ha='right', **single_kwargs)

    buf = io.BytesIO()
    fig.savefig(buf, format='png', dpi=dpi, bbox_inches='tight')
    plt.close(fig)
    buf.seek(0)
    data = buf.getvalue()
    buf.close()
    return data


def assemble_gif_from_bytes(frames_png_bytes: List[bytes], out_path: str, duration_ms: int):
    """Assemble PNG bytes (RGBA) into an animated GIF and save to out_path."""
    from PIL import Image as PILImage
    import io

    imgs = []
    for b in frames_png_bytes:
        buf = io.BytesIO(b)
        img = PILImage.open(buf).convert('RGBA')
        imgs.append(img.copy())
        buf.close()

    if not imgs:
        raise RuntimeError('No frames to assemble')

    outp = Path(out_path)
    outp.parent.mkdir(parents=True, exist_ok=True)
    imgs[0].save(outp, save_all=True, append_images=imgs[1:], duration=duration_ms, loop=0, optimize=False, disposal=2)
